#!/usr/bin/env python3
"""
强化版Bark模型下载器 - 防睡眠中断
"""

import os
import sys
import time
import subprocess
import signal
from pathlib import Path

def prevent_sleep():
    """防止系统睡眠"""
    try:
        # 使用caffeinate防止系统睡眠
        process = subprocess.Popen(['caffeinate', '-d', '-i', '-s'])
        return process
    except Exception as e:
        print(f"警告: 无法启动防睡眠: {e}")
        return None

def download_with_retry(max_retries=3):
    """带重试机制的下载"""
    
    # 设置环境变量
    os.environ['XDG_CACHE_HOME'] = os.path.expanduser('~/.cache')
    os.environ['XDG_DATA_HOME'] = os.path.expanduser('~/.local/share')
    os.environ['TTS_CACHE_DIR'] = os.path.expanduser('~/.cache/tts')
    
    try:
        from TTS.api import TTS
        print("✅ TTS库导入成功")
    except ImportError as e:
        print(f"❌ TTS库导入失败: {e}")
        return False
    
    for attempt in range(max_retries):
        try:
            print(f"\n🔄 尝试下载 ({attempt + 1}/{max_retries})...")
            print("📍 下载Bark模型，请保持网络连接...")
            
            # 强制重新下载
            model_path = Path.home() / ".local/share/tts/tts_models--multilingual--multi-dataset--bark"
            if model_path.exists():
                import shutil
                shutil.rmtree(model_path)
                print("🗑️ 清理旧模型文件")
            
            # 下载模型
            tts = TTS(model_name="tts_models/multilingual/multi-dataset/bark", progress_bar=True)
            
            # 测试模型
            print("🧪 测试模型...")
            test_text = "Hello, this is a test."
            tts.tts_to_file(
                text=test_text,
                file_path="/tmp/test_bark_robust.wav"
            )
            
            # 清理测试文件
            test_file = Path("/tmp/test_bark_robust.wav")
            if test_file.exists():
                test_file.unlink()
            
            print("✅ Bark模型下载并测试成功！")
            return True
            
        except KeyboardInterrupt:
            print("\n🛑 用户中断下载")
            return False
            
        except Exception as e:
            print(f"❌ 下载失败 (尝试 {attempt + 1}): {e}")
            if attempt < max_retries - 1:
                wait_time = (attempt + 1) * 30  # 递增等待时间
                print(f"⏳ {wait_time}秒后重试...")
                time.sleep(wait_time)
            else:
                print("💥 所有重试都失败了")
                return False
    
    return False

def main():
    print("="*60)
    print("🎭 强化版Bark模型下载器")
    print("🛡️ 防睡眠中断 + 自动重试")
    print("="*60)
    
    # 启动防睡眠
    caffeinate_process = prevent_sleep()
    
    try:
        # 执行下载
        success = download_with_retry()
        
        if success:
            print("\n🎉 Bark模型下载完成！")
            print("📝 更新状态文件...")
            
            # 更新状态文件
            status_content = f"""# Bark TTS模型下载状态记录

## 当前时间: {time.strftime('%Y-%m-%d %H:%M')}

## 下载状态: ✅ 完成

## 已完成下载的文件:
- ✅ 所有Bark模型文件下载完成
- ✅ 模型测试通过

## 文件位置:
- 主目录: `{Path.home()}/.local/share/tts/tts_models--multilingual--multi-dataset--bark/`

## 状态:
- Bark TTS功能已就绪
- 可以正常使用多人对话功能
"""
            
            with open("bark_download_status.md", "w", encoding="utf-8") as f:
                f.write(status_content)
                
            print("✅ 状态文件已更新")
            return True
        else:
            print("\n💥 下载失败")
            return False
            
    finally:
        # 停止防睡眠进程
        if caffeinate_process:
            try:
                caffeinate_process.terminate()
                caffeinate_process.wait(timeout=5)
                print("🔓 防睡眠已停止")
            except:
                caffeinate_process.kill()

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)